# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: leon

import logging
from typing import List, Union

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase


class UniswapV3WsSubscriber(WebsocketJsonSubscriberBase):
    def __init__(self,
                 queue,
                 group_id: str,
                 worker_id: str,
                 symbol: Union[List[str], str],
                 provider: str = 'ws://52.74.155.22:9444/ws',
                 blockchain_name: str = 'ethereum',
                 close_callback=None,
                 logger=None):
        super().__init__(close_callback, logger or logging.getLogger(__name__))
        self._symbol = to_list(symbol)

        self._register_url(provider)
        self._blockchain_name = blockchain_name
        self._publisher = None
        self._init_publisher(queue, group_id, worker_id)

    def _init_publisher(self, queue, group_id: str, worker_id: str):
        if self._blockchain_name == 'ethereum':
            topic_expect_bbo = topic_map.search_topic_id(
                'uniswapv3_ws_depth_%s%s' % (group_id, worker_id))[0]
            topic_trade = topic_map.search_topic_id(
                'uniswapv3_ws_trade_%s%s' % (group_id, worker_id))[0]
        elif self._blockchain_name == 'arbitrum':
            topic_expect_bbo = topic_map.search_topic_id(
                'uniswapv3-arbitrum_ws_depth_%s%s' % (group_id, worker_id))[0]
            topic_trade = topic_map.search_topic_id(
                'uniswapv3-arbitrum_ws_trade_%s%s' % (group_id, worker_id))[0]
        else:
            raise NotImplementedError(self._blockchain_name)

        self._publisher = {
            'depth': queue.get_writer(topic_id=topic_expect_bbo[1]),
            'trade': queue.get_writer(topic_id=topic_trade[1]),
        }

        self._enable_topic_event_publisher(
            queue,
            'uniswapv3-arbitrum_rest_topic_event',
            [
               topic_expect_bbo[0],
               topic_trade[0],
            ])
        self._add_teardown_callback(self._teardown_publisher)

    def _teardown_publisher(self):
        if self._publisher:
            for _, publisher in self._publisher.items():
                publisher.close()
            self._publisher = None

    def on_message(self, timestamp: int, message: dict):
        if self.stopped:
            return
        event_type = message.get('topic', None)
        try:
            self._publisher[event_type].write(timestamp, message)
        except KeyError:
            self._logger.error('Unknown event_type: %s' % event_type)
